Write a custom CUDA kernel to optimize Generalized Cross Entropy (GCE) Loss.

Formula: Loss = (1 - Pt^q) / q
Where Pt is the softmax probability of the target class: exp(logit_t) / sum(exp(logits)).
This loss is robust to noisy labels (NeurIPS 2018).

Problem Analysis:
1. Memory Bottleneck: The standard implementation F.softmax(logits) -> gather(targets) -> pow -> div creates a full-sized (N, C) probability tensor. This is wasteful as only the probability of the target class is needed for the loss.
2. Bandwidth Waste: Writing and reading the intermediate Softmax tensor consumes significant global memory bandwidth.

Optimization Strategy: Fused Softmax-GCE Kernel

The goal is to perform Softmax normalization and GCE calculation in a single pass without materializing the probability matrix.

1. Row-wise Parallelism: Assign one CUDA Block to process one sample (row) of the logits.

2. Fused Reduction:
Pass 1 (Max): Compute the maximum logit value M in the row for numerical stability.
Pass 2 (Sum): Compute the sum of exponentials S = sum(exp(x_i - M)).
Target Extraction: During iteration, identify and store the logit value corresponding to the target label logit_t.

3. Vectorized Access: Use float4 loads to maximize memory throughput when reading logits.

4. In-Register Calculation:
Compute Pt = exp(logit_t - M) / S.
Compute Loss = (1 - pow(Pt, q)) / q.
Write the single scalar loss to global memory.

This approach reduces global memory writes by a factor of C (Classes).
  
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:   
  
```python
import torch
import torch.nn as nn
import torch.nn.functional as F

BATCH_SIZE = 2048
NUM_CLASSES = 4096
SHAPE = (BATCH_SIZE, NUM_CLASSES)

Q_VALUE = 0.7
REDUCTION = 'none'

class GCELoss(nn.Module):
    """
    Generalized Cross Entropy Loss
    L = (1 - Pt^q) / q
    """
    def __init__(self, q=0.7, reduction='mean'):
        super(GCELoss, self).__init__()
        self.q = q
        self.reduction = reduction
        self.epsilon = 1e-7

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        probs = F.softmax(logits, dim=1)
        probs_t = probs.gather(1, targets.unsqueeze(1)).squeeze(1)
        
        probs_t = probs_t.clamp(min=self.epsilon, max=1.0)
        loss = (1.0 - probs_t.pow(self.q)) / self.q
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, q=0.7, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = GCELoss(q=q, reduction=reduction)
    
    def forward(self, logits, targets):
        return self.loss_fn(logits, targets)

def get_inputs():
    logits = torch.randn(SHAPE, dtype=torch.float32)
    targets = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
    return [logits.contiguous(), targets.contiguous()]

def get_init_inputs():
    return [Q_VALUE, REDUCTION]